import torch
import cv2
from PIL import Image
from torchvision import transforms
import numpy as np
from resnet import resNet
import glob

# GPU CPU选择
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

net = resNet()
net.to(device)

net.load_state_dict(torch.load("E:\pythonProject\project1\models\model_resnet"))

im_test = glob.glob("E:\pythonProject\dataset\cifar10\cifar-10-batches-py\TEST\*\*")

np.random.shuffle(im_test)

label_name = ["airplane",
              "automobile",
              "bird",
              "cat",
              "deer",
              "dog",
              "frog",
              "horse",
              "ship",
              "truck"]

test_transform = transforms.Compose([
    transforms.CenterCrop((28, 28)),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

for im_path in im_test:
    net.eval()
    im_data = Image.open(im_path).convert("RGB")
    inputs = test_transform(im_data)
    inputs = torch.unsqueeze(inputs, dim=0)
    inputs = inputs.to(device)
    outputs = net(inputs)

    _, pred = torch.max(outputs.data, dim=1)
    print(label_name[pred.cpu().numpy()[0]])

    img = np.array(im_data)
    img = img[:, :, [1, 2, 0]]

    img = cv2.resize(img, (300, 300))
    cv2.imshow("im", img)
    cv2.waitKey()
